{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run upon export from spreadsheet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import random\n",
    "import time\n",
    "\n",
    "from astropy.io import fits\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "pd.set_option('display.max_columns', None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "final = pd.DataFrame()\n",
    "\n",
    "labels = pd.read_csv('/mnt/tess/labels/labels-triage-v13.csv', header=0, low_memory=False)\n",
    "labels = labels[['TIC ID', 'Final', 'av', 'md', 'ch', 'as', 'mk', 'et', 'dm', 'td']]\n",
    "raters = ['av', 'md', 'ch', 'as', 'mk', 'et', 'dm', 'td']\n",
    "\n",
    "tr6 = pd.read_csv('/mnt/tess/astronet/tces-v6-train.csv', header=0, low_memory=False)\n",
    "tr7 = pd.read_csv('/mnt/tess/astronet/tces-v7-train.csv', header=0, low_memory=False)\n",
    "vl6 = pd.read_csv('/mnt/tess/astronet/tces-v6-val.csv', header=0, low_memory=False)\n",
    "vl7 = pd.read_csv('/mnt/tess/astronet/tces-v7-val.csv', header=0, low_memory=False)\n",
    "ts6 = pd.read_csv('/mnt/tess/astronet/tces-v6-test.csv', header=0, low_memory=False)\n",
    "ts7 = pd.read_csv('/mnt/tess/astronet/tces-v7-test.csv', header=0, low_memory=False)\n",
    "\n",
    "origtce = pd.read_csv('/mnt/tess/labels/tce_bls_instar.csv', header=0, low_memory=False)\n",
    "origtce['disp_E'] = 0\n",
    "origtce['disp_S'] = 0\n",
    "origtce['disp_B'] = 0\n",
    "origtce['disp_J'] = 0\n",
    "origtce['disp_N'] = 0\n",
    "origtce['Duration'] /= 24.0\n",
    "origtce2 = pd.read_csv('/mnt/tess/labels/tce_bls_instar+old.csv', header=0, low_memory=False)\n",
    "origtce2['disp_E'] = 0\n",
    "origtce2['disp_S'] = 0\n",
    "origtce2['disp_B'] = 0\n",
    "origtce2['disp_J'] = 0\n",
    "origtce2['disp_N'] = 0\n",
    "origtce2['Duration'] /= 24.0\n",
    "\n",
    "newtce = pd.read_csv('/mnt/tess/labels/tces-triage-v12.csv', header=0, low_memory=False)\n",
    "\n",
    "mast = pd.read_csv('/mnt/tess/labels/ext_mast_data.csv', header=0, low_memory=False)\n",
    "\n",
    "\n",
    "print('...', end='')\n",
    "minidx = 0\n",
    "\n",
    "for idx, rec in labels.iterrows():\n",
    "  ticid = rec['TIC ID']\n",
    "  if idx < minidx:\n",
    "    continue\n",
    "\n",
    "  rec = rec.copy()\n",
    "  tc = tr6[tr6['TIC ID'] == ticid]\n",
    "  tc7 = tr7[tr7['TIC ID'] == ticid]\n",
    "  split = 'train'\n",
    "  \n",
    "  if not len(tc):\n",
    "    tc = vl6[vl6['TIC ID'] == ticid]\n",
    "    tc7 = vl7[vl7['TIC ID'] == ticid]\n",
    "    split = 'val'\n",
    "  if not len(tc):\n",
    "    tc = ts6[ts6['TIC ID'] == ticid]\n",
    "    tc7 = ts7[ts7['TIC ID'] == ticid]\n",
    "    split = 'test'\n",
    "\n",
    "  obj_type = mast[mast['tic_id'] == ticid]['objType']\n",
    "  \n",
    "  isy3 = idx > 26380\n",
    "\n",
    "  if isy3:\n",
    "    tc = newtce[(newtce['tic_id'] == ticid)\n",
    "              & (newtce['filename'].str.startswith('mk_'))]\n",
    "    assert len(tc) == 1, ticid\n",
    "    split = tc['Split'].values[0]\n",
    "    assert split in ('train', 'test', 'val')\n",
    "  else:\n",
    "    if len(tc):\n",
    "      assert len(tc) == 1, ticid  \n",
    "      assert len(tc7) == 1, ticid\n",
    "      assert tc['RA'].values == tc7['RA'].values\n",
    "      assert tc['Dec'].values == tc7['Dec'].values\n",
    "      assert tc['Tmag'].values == tc7['Tmag'].values\n",
    "      assert tc['Epoc'].values == tc7['Epoc'].values\n",
    "      assert tc['Period'].values == tc7['Period'].values\n",
    "      assert tc['Duration'].values == tc7['Duration'].values\n",
    "      assert tc['Transit_Depth'].values == tc7['Transit_Depth'].values\n",
    "      assert (tc['star_rad'].isna().values == tc7['star_rad'].isna().values) or tc['star_rad'].values == tc7['star_rad'].values\n",
    "      assert (tc['star_mass'].isna().values == tc7['star_mass'].isna().values) or tc['star_mass'].values == tc7['star_mass'].values\n",
    "      assert (tc['teff'].isna().values == tc7['teff'].isna().values) or tc['teff'].values == tc7['teff'].values\n",
    "      assert (tc['logg'].isna().values == tc7['logg'].isna().values) or tc['logg'].values == tc7['logg'].values\n",
    "      assert tc['SN'].values == tc7['SN'].values\n",
    "      assert tc['Qingress'].values == tc7['Qingress'].values\n",
    "      assert (len(obj_type) == 0) or (obj_type.values[0] == 'STAR')\n",
    "    else:\n",
    "      if len(obj_type.values) and obj_type.values[0] == 'STAR':\n",
    "        tc = origtce[origtce['tic_id'] == ticid]\n",
    "        if not len(tc):\n",
    "          tc = origtce2[origtce2['tic_id'] == ticid]\n",
    "        assert len(tc), ticid\n",
    "        assert len(tc['Ilabel'].values), ticid\n",
    "        instar = tc['Ilabel'].values[0]\n",
    "        if instar:\n",
    "          split = 'instar'\n",
    "        else:\n",
    "          if rec['Final'] == 'U':\n",
    "            split = 'instar'\n",
    "          else:\n",
    "            split = random.choice(['train', 'teat', 'val'])\n",
    "      else:\n",
    "        tc = origtce[origtce['tic_id'] == ticid]\n",
    "        if not len(tc):\n",
    "          tc = origtce2[origtce2['tic_id'] == ticid]\n",
    "        split = 'notstar'\n",
    "        assert len(tc), ticid\n",
    "    assert 'filename' not in tc\n",
    "\n",
    "  _, tc = next(tc.iterrows())\n",
    "  \n",
    "  rec['RA'] = tc['RA']\n",
    "  rec['Dec'] = tc['Dec']\n",
    "  rec['Tmag'] = tc['Tmag']\n",
    "  rec['Epoc'] = tc['Epoc']\n",
    "  rec['Per'] = tc['Period']\n",
    "  rec['Dur'] = tc['Duration']\n",
    "  rec['Depth'] = tc['Transit_Depth']\n",
    "  rec['SRad'] = tc['star_rad']\n",
    "  rec['SMass'] = tc['star_mass']\n",
    "  \n",
    "  rec['split'] = split\n",
    "  \n",
    "  if not isy3:\n",
    "    if tc['disp_E'] > 0:\n",
    "      assert rec['Final'] == 'E' or (\n",
    "        ticid in (293346794,  # Changed later\n",
    "                  294871044,  # Changed later\n",
    "                  92328347,  # Changed later\n",
    "                  83408987,  # Changed later\n",
    "                  268608949,  # Changed later\n",
    "                 )\n",
    "        ), ticid\n",
    "      assert tc['disp_S'] == 0\n",
    "      assert tc['disp_B'] == 0\n",
    "      assert tc['disp_J'] == 0\n",
    "      assert tc['disp_N'] == 0\n",
    "    elif tc['disp_S'] > 0:\n",
    "      assert rec['Final'] == 'S'\n",
    "      assert tc['disp_E'] == 0\n",
    "      assert tc['disp_B'] == 0\n",
    "      assert tc['disp_J'] == 0\n",
    "      assert tc['disp_N'] == 0\n",
    "    else:\n",
    "      if tc['disp_B'] > 0:\n",
    "        if rec['Final'] == 'B':\n",
    "          assert tc['disp_E'] == 0\n",
    "          assert tc['disp_S'] == 0\n",
    "          assert tc['disp_J'] == 0\n",
    "          assert tc['disp_N'] == 0\n",
    "        else:\n",
    "          assert np.isnan(rec['Final'])\n",
    "          assert tc['disp_B'] == sum(rec[r] == 'B' for r in raters)\n",
    "      if tc['disp_J'] > 0:\n",
    "        if rec['Final'] == 'J':\n",
    "          assert tc['disp_E'] == 0\n",
    "          assert tc['disp_S'] == 0\n",
    "          assert tc['disp_B'] == 0\n",
    "          assert tc['disp_N'] == 0\n",
    "        else:\n",
    "          assert np.isnan(rec['Final'])\n",
    "          assert tc['disp_J'] == sum(rec[r] == 'J' for r in raters), (idx, ticid)\n",
    "      if tc['disp_N'] > 0:\n",
    "        if rec['Final'] == 'N':\n",
    "          assert tc['disp_E'] == 0\n",
    "          assert tc['disp_S'] == 0\n",
    "          assert tc['disp_B'] == 0\n",
    "          assert tc['disp_J'] == 0\n",
    "        else:\n",
    "          assert np.isnan(rec['Final'])\n",
    "          assert tc['disp_N'] == sum(rec[r] == 'N' for r in raters)\n",
    "        \n",
    "  prefs = {\n",
    "    339672028: 's0013',\n",
    "    425721385: 's0013',\n",
    "    300557619: 's0012',\n",
    "    327301957: 's0013',\n",
    "    55559618:  's0012',\n",
    "    351603103: 's0013',\n",
    "    278866211: 's0012',\n",
    "    391821647: 's0008',\n",
    "    281979481: 's0013',\n",
    "    314865962: 's0012',\n",
    "    394722182: 's0013',\n",
    "    383390264: 's0013',\n",
    "    287156968: 's0013',\n",
    "    167418903: 's0013',\n",
    "    261257684: 's0013',\n",
    "    311890977: 's0013',\n",
    "    366989877: 's0013',\n",
    "    38571020: 's0013',\n",
    "    280097543: 's0013',\n",
    "    373844472: 's0013',\n",
    "    234345288: 's0013',\n",
    "    261136679: 's0013',\n",
    "    469782185: 's0013',\n",
    "    300871545: 's0012',\n",
    "    260304296: 's0012',\n",
    "    401604346: 's0013',\n",
    "    269450900: 's0013',\n",
    "    295541511: 's0013',\n",
    "    238197709: 's0012',\n",
    "    421894914: 's0013',\n",
    "    446549905: 's0013',\n",
    "    6354567: '2019247',\n",
    "    55452495: 's0013',\n",
    "    263003176: 's0013',\n",
    "    318608749: 's0013',\n",
    "    253990973: 's0013',\n",
    "    406976746: 's0013',\n",
    "    278683844: 's0012',\n",
    "    299799658: 's0013',\n",
    "    271596418: 's0012',\n",
    "    350153977: 's0013',\n",
    "    143257766: 's0013',\n",
    "    91576611: 's0013',\n",
    "    300810086: 's0012',\n",
    "    177077336: 's0013',\n",
    "    278415929: 's0013',\n",
    "    388130235: 's0013',\n",
    "    425561347: 's0013',\n",
    "    140760434: 's0013',\n",
    "    382625239: 's0013',\n",
    "    270341214: 's0013',\n",
    "    41227743: 's0013',\n",
    "    304100538: 's0013',\n",
    "    179317684: 's0012',\n",
    "    29831208: 's0012',\n",
    "    348770361: 's0013',\n",
    "    271900960: 's0012',\n",
    "    370133522: 's0013',\n",
    "    260985861: 's0013',\n",
    "    375059587: 's0012',\n",
    "    299780329: 's0013',\n",
    "    380783252: 's0013',\n",
    "    360742636: 's0013',\n",
    "    287329267: 's0012',\n",
    "    277099925: 's0012',\n",
    "    325468685: 's0013',\n",
    "    280095254: 's0013',\n",
    "    382626661: 's0012',\n",
    "    348844154: 's0012',\n",
    "    339733013: 's0012',\n",
    "    150428135: 's0013',\n",
    "    261369656: 's0013',\n",
    "    281924357: 's0012',\n",
    "    409934330: 's0013',\n",
    "    300293197: 's0013',\n",
    "    372172128: 's0013',\n",
    "    349373192: 's0012',\n",
    "    271581073: 's0013',\n",
    "    379286801: 's0013',\n",
    "    306263608: '202107',\n",
    "    155001079: 's0017',\n",
    "    267489265: '2020262000000-s0016',\n",
    "    116608612: '202107',\n",
    "    417931607: 's0022',\n",
    "    125735470: '202107',\n",
    "    199376584: '2021070000000-s0015',\n",
    "    188768068: 's0017',\n",
    "    312083267: '2021070000000-s0015',\n",
    "    387664866: '2021070000000-s0016',\n",
    "    387834907: 's0025',\n",
    "    277566483: '2021070000000-s0016',\n",
    "    27491137: 's0023',\n",
    "    375506058: '2021070000000-s0016',\n",
    "    142387023: 's0022',\n",
    "    115771549: '2021070000000-s0016',\n",
    "    21755546: '2021070000000-s0014',\n",
    "    69679391: '2021070000000-s0014',\n",
    "    219857012: 's0026',\n",
    "    136971594: '2021070000000-s0016',\n",
    "    91987762: '2021070000000-s0021',\n",
    "    353782445: 's0026',\n",
    "    298663873: 's0026',\n",
    "    172518755: '2021070000000-s0021',\n",
    "    390651552: '2021070000000-s0023',\n",
    "    352764091: 's0024',\n",
    "    198384408: 's0026',\n",
    "    1715469662: 's0026',\n",
    "    138588540: 's0021',\n",
    "    129979528: '2021070000000-s0018',\n",
    "    311121985: '2021070000000-s0019',\n",
    "    198213332: 's0026',\n",
    "    207468071: 's0025',\n",
    "    180695581: 's0023',\n",
    "    356235833: '2021070000000-s0026',\n",
    "    288132261: 's0026',\n",
    "    287080092: 's0026',\n",
    "    117789567: '2021070000000-s0014',\n",
    "    350995812: '2021070000000-s0014',\n",
    "    120602501: '2021070000000-s0016',\n",
    "    356700488: 's0026',\n",
    "    198187049: 's0026',\n",
    "    284450803: 's0026',\n",
    "    160268701: 's0021',\n",
    "    142276270: 's0022',\n",
    "    286923464: 's0022',\n",
    "    21960306: 's0026',\n",
    "    202563254: '2021070000000-s0015',\n",
    "    118301361: '2021070000000-s0018',\n",
    "    116483514: 's0024',\n",
    "    298647682: 's0025',\n",
    "    365683032: '2021070000000-s0016',\n",
    "    365683032: '2021070000000-s0016',\n",
    "    264171144: 's0026',\n",
    "    368435330: '2021070000000-s0022',\n",
    "    138819293: '2021070000000-s0022',\n",
    "    233684293: '2021070000000-s0021',\n",
    "    122220263: '2021070000000-s0014',\n",
    "    1400770435: 's0026',\n",
    "    202426247: '2021070000000-s0023',\n",
    "    459978312: 's0026',\n",
    "    269701147: 's0025',\n",
    "    349827430: 's0021',\n",
    "    192372961: '2021070000000-s0018',\n",
    "    367366318: '2021070000000-s0019',\n",
    "    26880783: '2021070000000-s0016',\n",
    "    155114483: 's0017',\n",
    "    275111900: '2021070000000-s0016',\n",
    "    367607434: 's0017',\n",
    "    137284119: '2021070000000-s0016',\n",
    "    2676497: 's0019',\n",
    "    395171208: '2021070000000-s0016',\n",
    "    129539786: '2021070000000-s0016',\n",
    "    377191482: '2021070000000-s0017',\n",
    "    406672232: '2021070000000-s0014',\n",
    "    354006740: '2021070000000-s0019',\n",
    "    233009109: 's0026',\n",
    "    237184773: 's0026',\n",
    "    82308728: '2021070000000-s0022',\n",
    "    198356533: 's0026',\n",
    "    364186197: 's0025',\n",
    "    69997672: '2021070000000-s0014',\n",
    "    320525204: 's0026',\n",
    "    298073824: '2021070000000-s0015',\n",
    "    110428269: '2021070000000-s0014',\n",
    "    89389197: '2021070000000-s0014',\n",
    "    350132371: 's0025',\n",
    "    219751469: 's0026',\n",
    "    26547036: '2021070000000-s0016',\n",
    "    28230919: '2021070000000-s0015',\n",
    "    138017750: '2021070000000-s0018',\n",
    "    293954617: '2021070000000-s0016',\n",
    "    66445643: 's0017',\n",
    "    224600500: 's0026',\n",
    "    159510109: 's0026',\n",
    "    16740101: '2021070000000-s0015',\n",
    "    237222864: 's0026',\n",
    "    198457103: 's0026',\n",
    "    229455001: 's0026',\n",
    "  }\n",
    "#   minidx = 22470\n",
    "  if isy3:\n",
    "    rec['Year'] = 3\n",
    "    rec['File'] = tc['filename']\n",
    "    f = fits.open('/mnt/tess/lc/' + tc['filename'])\n",
    "    rec['MinT'] = min(f[1].data[\"TIME\"])\n",
    "    rec['MaxT'] = max(f[1].data[\"TIME\"])\n",
    "  else:\n",
    "    rec['Year'] = 2\n",
    "    if (rec['Final'] == 'U'):\n",
    "      rec['File'] = ''\n",
    "      rec['MinT'] = -1\n",
    "      rec['MaxT'] = -1\n",
    "      rec['split'] = 'norpt'\n",
    "    else:\n",
    "      if ticid in prefs:\n",
    "        pat = '/mnt/tess/lc/*' + prefs[ticid] + '*0' + str(ticid) + '*'\n",
    "      else:\n",
    "        pat = '/mnt/tess/lc/*0' + str(ticid) + '*'\n",
    "      lc = glob.glob(pat)\n",
    "      assert len(lc), pat\n",
    "      flc = []\n",
    "      for p in lc:\n",
    "        t = time.localtime(os.path.getmtime(p))\n",
    "        if (t.tm_year, t.tm_mon) < (2021, 8):\n",
    "          flc.append(p)\n",
    "      if len(flc) != 1:\n",
    "        print(ticid, pat)\n",
    "        print('\\n'.join(flc))\n",
    "        assert False\n",
    "      flc, = flc\n",
    "      rec['File'] = os.path.basename(flc)\n",
    "\n",
    "      f = fits.open(flc)\n",
    "      rec['MinT'] = min(f[1].data[\"TIME\"])\n",
    "      rec['MaxT'] = max(f[1].data[\"TIME\"])\n",
    "  \n",
    "  final = final.append(rec)\n",
    "  \n",
    "  print(f'\\r{idx}   ', end='')\n",
    "#   if idx > 5:\n",
    "#     break\n",
    "  \n",
    "final = final.astype({'Depth': int, 'TIC ID': int, 'Year': int})\n",
    "  \n",
    "final.sample()\n",
    "\n",
    "final.to_csv('aligned.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.set_option('display.max_columns', None)\n",
    "final"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
